[TRTLLM-11875][feat] Support fine-grained context chunk management by VALLIS-NERIA · Pull Request #16427 · NVIDIA/TensorRT-LLM

VALLIS-NERIA · 2026-07-15T13:15:49Z

Summary by CodeRabbit

New Features
- Added snapshot-aware context chunking to align processing with recurrent-state snapshot points.
- Exposed expected snapshot points through the Python request interface.
- Improved cache reuse scheduling across hybrid model execution.
Bug Fixes
- Improved chunk sizing under context and compute-capacity limits.
- Corrected KV-head selection when earlier layers report zero heads.
- Added validation for incompatible hybrid cache transceiver configurations.
- Improved cache buffer sizing and alignment for FlashMLA setups.

Description

This PR has two goals:

Fix several independent KV-cache correctness issues.
Add expect_snapshot_points as the manager-to-scheduler contract and move the existing C++ Mamba interval-reuse chunking onto that contract.

Snapshot-aware context chunking

Store expect_snapshot_points on the C++ request object and expose it through nanobind, so the C++ and Python schedulers share the same per-request field.
Have CppMambaHybridCacheManager publish regular states_snapshot_interval boundaries before scheduling when block reuse is enabled. Prompt end remains the implicit final boundary.
Make the C++ MicroBatchScheduler, Python PyMicroBatchScheduler, and KVCacheV2Scheduler honor these absolute boundaries under FORCE_CHUNK.
When no snapshot boundary remains, avoid proactive chunking and consume the remaining context; only budget and max-context limits may truncate it, with block-aligned truncation.
Add focused C++ and Python coverage for interval publication, successive boundaries, prompt-end handling, and capacity truncation.

KV-cache fixes

Calculate the default max_tokens_in_buffer only after the effective tokens_per_block is known, including the FlashMLA override, and round it up to a block boundary.
Publish the first non-zero per-layer KV-head count in hybrid-model RankInfo, because layer 0 may be a linear-attention layer with no KV heads.
Reject MixedMambaHybridCacheManager with a non-Python transceiver runtime, since its separately managed attention-KV and Mamba-state pools require the Python routing path.

This PR does not add the V2 Mamba cache manager, save-last reuse, or V2 Mamba disaggregated transfer.

Test Coverage

Full native rebuild for SM100 and SM103
C++ ForceChunkTest.*: 15 passed
Python TestForceChunkPolicy: 14 passed
Focused interval, KV-cache buffer, hybrid RankInfo, and runtime-guard tests: 13 passed
Mixed Mamba Python transceiver coverage across all 8 dtype combinations, plus runtime guards: 10 passed
Native request snapshot-point default, assignment, and copy probe
Complete pre-commit hook set

PR Checklist

Please check this after reviewing the above items as appropriate for this PR.

VALLIS-NERIA · 2026-07-15T13:16:09Z

/bot run

tensorrt-cicd · 2026-07-15T13:22:30Z

PR_Github #59453 [ run ] triggered by Bot. Commit: 7014727 Link to invocation

tensorrt-cicd · 2026-07-15T15:07:43Z

PR_Github #59453 [ run ] completed with state SUCCESS. Commit: 7014727
/LLM/main/L0_MergeRequest_PR pipeline #47921 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA

Author notes explaining the purpose and scope of each production change block and its focused regression coverage.

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-07-16T07:56:44Z

/bot run

tensorrt-cicd · 2026-07-16T08:02:50Z

PR_Github #59667 [ run ] triggered by Bot. Commit: 5a5b53d Link to invocation

tensorrt-cicd · 2026-07-16T08:19:04Z

PR_Github #59667 [ run ] completed with state FAILURE. Commit: 5a5b53d
/LLM/main/L0_MergeRequest_PR pipeline #48100 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-07-19T11:03:12Z

Updated the forced-chunk semantics in 76d8903, with the public-API-file cleanup in 31eed67:

Removed the redundant _force_chunk_size wrapper and the separate _next_expected_snapshot_point helper.
_get_forced_context_chunk_size now has one responsibility: choose the next pending absolute snapshot point, or return the full context_remaining_length when none remains.
Aligned PyMicroBatchScheduler, KVCacheV2Scheduler, and the C++ MicroBatchScheduler so chunk_unit_size is used only to align budget/max-context truncation, not as a proactive boundary.
Updated scheduler docs/code comments and C++/Python FORCE_CHUNK tests for the new no-boundary behavior.
Added a regression test confirming the normal executor path prepares snapshot points before scheduling, complementing the existing non-root PP path test.
Kept the documentation change out of llm_args.py, since this is an internal scheduler behavior and touching that public API file would unnecessarily require an API-compatibility label.

Validation:

21 targeted Python scheduler/executor tests passed.
Full changed-file pre-commit suite passed.

The PR description has also been updated to reflect that requests without a pending snapshot point consume the remaining context unless scheduling limits require truncation.

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-07-19T11:04:41Z

/bot run

tensorrt-cicd · 2026-07-19T11:10:41Z

PR_Github #60171 [ run ] triggered by Bot. Commit: 31eed67 Link to invocation

tensorrt-cicd · 2026-07-19T11:40:03Z

PR_Github #60171 [ run ] completed with state FAILURE. Commit: 31eed67
/LLM/main/L0_MergeRequest_PR pipeline #48546 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA · 2026-07-20T02:44:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T02:50:35Z

PR_Github #60226 [ run ] triggered by Bot. Commit: 383894d Link to invocation

tensorrt-cicd · 2026-07-20T06:35:01Z

PR_Github #60226 [ run ] completed with state SUCCESS. Commit: 383894d
/LLM/main/L0_MergeRequest_PR pipeline #48593 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA · 2026-07-20T06:41:39Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T06:47:09Z

PR_Github #60298 [ run ] triggered by Bot. Commit: 383894d Link to invocation

tensorrt-cicd · 2026-07-20T11:37:48Z

PR_Github #60298 [ run ] completed with state SUCCESS. Commit: 383894d
/LLM/main/L0_MergeRequest_PR pipeline #48649 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…ent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-07-20T12:33:56Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T12:39:36Z

PR_Github #60370 [ run ] triggered by Bot. Commit: 6797080 Link to invocation

tensorrt-cicd · 2026-07-20T16:17:45Z

PR_Github #60370 [ run ] completed with state FAILURE. Commit: 6797080
/LLM/main/L0_MergeRequest_PR pipeline #48709 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA · 2026-07-20T16:25:10Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T16:34:24Z

PR_Github #60407 [ run ] triggered by Bot. Commit: 6797080 Link to invocation

tensorrt-cicd · 2026-07-20T18:17:18Z

PR_Github #60407 [ run ] completed with state FAILURE. Commit: 6797080
/LLM/main/L0_MergeRequest_PR pipeline #48747 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA · 2026-07-20T18:19:41Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T18:28:18Z

PR_Github #60431 [ run ] triggered by Bot. Commit: 6797080 Link to invocation

tensorrt-cicd · 2026-07-21T00:00:16Z

PR_Github #60431 [ run ] completed with state SUCCESS. Commit: 6797080
/LLM/main/L0_MergeRequest_PR pipeline #48767 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA · 2026-07-21T00:03:41Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-21T00:10:38Z

PR_Github #60504 [ run ] triggered by Bot. Commit: 6797080 Link to invocation

tensorrt-cicd · 2026-07-21T02:40:31Z

PR_Github #60504 [ run ] completed with state SUCCESS. Commit: 6797080
/LLM/main/L0_MergeRequest_PR pipeline #48829 completed with status: 'SUCCESS'

CI Report

Link to invocation

github-actions Bot assigned VALLIS-NERIA Jul 15, 2026

VALLIS-NERIA commented Jul 16, 2026

View reviewed changes

VALLIS-NERIA mentioned this pull request Jul 16, 2026

[None][feat] Add V2 Mamba hybrid cache with snapshot reuse #16471

Closed

1 task

VALLIS-NERIA added 7 commits July 16, 2026 15:07

[None][fix] Harden scheduling and disagg lifecycle (By Agent)

f2115d5

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

[None][chore] Apply resource manager test formatting (By Agent)

16e608e

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

[None][fix] Address scheduling and transfer review (By Agent)

3456ec3

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

[None][fix] Align hybrid Mamba C++ disagg routing (By Agent)

75651dd

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

[None][test] Wait for hybrid transfer cancellation (By Agent)

fb13a38

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

[None][test] Cover mixed Mamba Python transfer (By Agent)

47e8a7b

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

[None][test] Cover hybrid runtime auto fallback (By Agent)

5a5b53d

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA force-pushed the agent/mamba-scheduling-disagg-lifecycle branch from 7014727 to 5a5b53d Compare July 16, 2026 07:42

VALLIS-NERIA changed the title ~~[None][fix] Harden scheduling and disaggregated transfer lifecycle~~ [None][feat] Support fine-grained context chunk management Jul 16, 2026

VALLIS-NERIA mentioned this pull request Jul 16, 2026

[None][fix] Require Python transceiver for mixed Mamba #16428

Closed

1 task

VALLIS-NERIA mentioned this pull request Jul 16, 2026

[None][feat] Add V2 Mamba hybrid cache with snapshot reuse VALLIS-NERIA/TensorRT-LLM#8

Draft

1 task

[None][fix] Address PR1 review feedback (By Agent)

d935140

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA commented Jul 16, 2026

View reviewed changes

VALLIS-NERIA marked this pull request as ready for review July 16, 2026 12:25

VALLIS-NERIA requested review from a team as code owners July 16, 2026 12:25

VALLIS-NERIA requested review from JunyiXu-nv, Tabrizian and zhaoyangwang-nvidia July 16, 2026 12:25

[Agent fix] Keep forced-chunk docs internal (By Agent)

31eed67

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

litaotju approved these changes Jul 19, 2026

View reviewed changes

Merge branch 'main' into agent/mamba-scheduling-disagg-lifecycle

383894d

QiJune reviewed Jul 20, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py Outdated

QiJune approved these changes Jul 20, 2026

View reviewed changes

[None][fix] Preserve attention capacity on Mamba-free PP ranks (By Ag…

6797080

…ent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA enabled auto-merge (squash) July 20, 2026 23:50

VALLIS-NERIA merged commit bb90835 into NVIDIA:main Jul 21, 2026
8 checks passed

Uh oh!

Conversation

VALLIS-NERIA commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Snapshot-aware context chunking

KV-cache fixes

Test Coverage

PR Checklist

Uh oh!

VALLIS-NERIA commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

VALLIS-NERIA left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

VALLIS-NERIA commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

VALLIS-NERIA commented Jul 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

VALLIS-NERIA commented Jul 19, 2026

Uh oh!

tensorrt-cicd commented Jul 19, 2026

Uh oh!

tensorrt-cicd commented Jul 19, 2026

Uh oh!

VALLIS-NERIA commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

VALLIS-NERIA commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

VALLIS-NERIA commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

VALLIS-NERIA commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

VALLIS-NERIA commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

VALLIS-NERIA commented Jul 15, 2026 •

edited

Loading

VALLIS-NERIA commented Jul 19, 2026 •

edited

Loading